from stock.utils.buffer_utils import get_stock_data
import os
from datetime import datetime
import torch
import numpy as np
from sklearn.preprocessing import MinMaxScaler
import os
import json
from datetime import datetime

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import TimeSeriesSplit
import matplotlib.pyplot as plt
import os
from copy import deepcopy


# 1. 数据处理器（直接预测收盘价）
class StockDataProcessor:
    def __init__(self, lookback=3):
        self.lookback = lookback
        self.scaler = StandardScaler()
        self._is_fitted = False

    def create_sequences(self, stock_a, stock_b):
        """创建输入序列和次日收盘价目标"""
        stock_a = np.array(stock_a, dtype=np.float32).reshape(-1, 1)
        stock_b = np.array(stock_b, dtype=np.float32).reshape(-1, 1)

        # 合并特征并标准化
        combined = np.hstack([stock_a, stock_b])
        combined_norm = self.scaler.fit_transform(combined)
        self._is_fitted = True

        # 生成序列
        X, y = [], []
        for i in range(len(combined) - self.lookback):
            # 输入特征：A和B的lookback天数据
            seq = combined_norm[i:i + self.lookback].flatten()

            # 目标：B股次日收盘价（标准化后）
            target = combined_norm[i + self.lookback, 1]  # 只取B股

            X.append(seq)
            y.append(target)

        return np.array(X), np.array(y).reshape(-1, 1), {
            'mean_': self.scaler.mean_,
            'scale_': self.scaler.scale_,
            'lookback': self.lookback
        }

    def augment_data(self, X, y, noise_scale=0.05):
        """数据增强策略"""
        # 添加高斯噪声
        X_noise = X + np.random.normal(0, noise_scale, X.shape)
        y_noise = y.copy()

        # 时间反转
        X_reverse = X[:, ::-1].copy()
        y_reverse = y.copy()

        # 幅度缩放
        scale = 1 + np.random.uniform(-0.1, 0.1, (len(X), 1))
        X_scale = X * scale
        y_scale = y * scale

        return (
            np.vstack([X, X_noise, X_reverse, X_scale]),
            np.vstack([y, y_noise, y_reverse, y_scale])
        )


# 2. LSTM预测模型
class StockPricePredictor(nn.Module):
    def __init__(self, input_size=6, hidden_size=16):
        super().__init__()
        self.lstm = nn.LSTM(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=1,
            batch_first=True,
            dropout=0.2
        )
        self.fc = nn.Sequential(
            nn.Linear(hidden_size, 8),
            nn.ReLU(),
            nn.Linear(8, 1)  # 直接预测标准化后的价格
        )

    def forward(self, x):
        lstm_out, _ = self.lstm(x)
        return self.fc(lstm_out[:, -1, :])  # 取最后一个时间步


# 3. 训练器
class StockTrainer:
    def __init__(self, model, model_dir='saved_models'):
        self.model = model
        self.model_dir = model_dir
        os.makedirs(model_dir, exist_ok=True)

    def train(self, X_train, y_train, scaler_params, epochs=500, batch_size=8):
        device = torch.device('cpu')
        self.model.to(device)

        # 准备数据加载器
        dataset = torch.utils.data.TensorDataset(
            torch.FloatTensor(X_train).unsqueeze(1),
            torch.FloatTensor(y_train)
        )
        loader = torch.utils.data.DataLoader(
            dataset, batch_size=batch_size, shuffle=True
        )

        # 使用Huber损失函数
        criterion = nn.HuberLoss()
        optimizer = optim.AdamW(self.model.parameters(), lr=0.001, weight_decay=1e-4)
        scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=20)

        # 训练循环
        best_loss = float('inf')
        history = {'train_loss': []}

        for epoch in range(epochs):
            epoch_loss = 0
            self.model.train()

            for batch_x, batch_y in loader:
                batch_x, batch_y = batch_x.to(device), batch_y.to(device)

                optimizer.zero_grad()
                outputs = self.model(batch_x)
                loss = criterion(outputs, batch_y)
                loss.backward()
                nn.utils.clip_grad_norm_(self.model.parameters(), 0.5)
                optimizer.step()

                epoch_loss += loss.item()

            avg_loss = epoch_loss / len(loader)
            history['train_loss'].append(avg_loss)
            scheduler.step(avg_loss)

            # 保存最佳模型
            if avg_loss < best_loss:
                best_loss = avg_loss
                self._save_model(scaler_params)

            if epoch % 50 == 0:
                print(f"Epoch {epoch}: Loss={avg_loss:.4f}")

        self._plot_training(history)
        return history

    def _save_model(self, scaler_params):
        """保存模型和参数"""
        torch.save({
            'model_state': self.model.state_dict(),
            'scaler_params': scaler_params,
            'model_config': {
                'input_size': self.model.lstm.input_size,
                'hidden_size': self.model.lstm.hidden_size
            }
        }, os.path.join(self.model_dir, 'best_model.pth'))

    def _plot_training(self, history):
        plt.plot(history['train_loss'], label='Training Loss')
        plt.xlabel('Epoch')
        plt.ylabel('Loss')
        plt.title('Training Progress')
        plt.legend()
        plt.show()


# 4. 预测器
class StockPricePredictorWrapper:
    def __init__(self, model_path='saved_models'):
        self.device = torch.device('cpu')
        self.model, self.scaler_params = self._load_model(model_path)
        self.lookback = self.scaler_params['lookback']

    def _load_model(self, model_path):
        """加载模型和scaler参数"""
        checkpoint = torch.load(
            os.path.join(model_path, 'best_model.pth'),
            map_location=self.device
        )

        # 初始化模型结构
        config = checkpoint['model_config']
        model = StockPricePredictor(
            input_size=config['input_size'],
            hidden_size=config['hidden_size']
        )
        model.load_state_dict(checkpoint['model_state'])
        model.to(self.device)
        model.eval()

        return model, checkpoint['scaler_params']

    def denormalize(self, normalized_price):
        """将标准化价格反归一化"""
        return normalized_price * self.scaler_params['scale_'][1] + self.scaler_params['mean_'][1]

    def predict(self, stock_a_values, stock_b_values):
        """
        预测次日收盘价
        :param stock_a_values: 股票A最近lookback天的收盘价列表
        :param stock_b_values: 股票B最近lookback天的收盘价列表
        :return: 预测的次日B股收盘价
        """
        if len(stock_a_values) != self.lookback or len(stock_b_values) != self.lookback:
            raise ValueError(f"需要输入{self.lookback}天的历史数据")

        # 归一化输入
        raw_input = np.hstack([
            np.array(stock_a_values).reshape(-1, 1),
            np.array(stock_b_values).reshape(-1, 1)
        ])
        scaled_input = (raw_input - self.scaler_params['mean_']) / self.scaler_params['scale_']
        input_seq = scaled_input.flatten()

        # 预测
        with torch.no_grad():
            input_tensor = torch.FloatTensor(input_seq).unsqueeze(0).unsqueeze(1).to(self.device)
            predicted_norm = self.model(input_tensor).item()

        # 反归一化
        return self.denormalize(predicted_norm)


# 示例使用
if __name__ == "__main__":
    # 模拟数据
    code = '101.HG00Y'
    list1 = get_stock_data(code)
    code = '1.600362'
    list2 = get_stock_data(code)
    stock_a = [x.close for x in list1]
    stock_b = [x.close for x in list2]
    # 1. 数据准备
    processor = StockDataProcessor(lookback=3)
    X, y, scaler_params = processor.create_sequences(stock_a, stock_b)
    X_aug, y_aug = processor.augment_data(X, y)
    print(f"原始数据量: {len(X)} | 增强后: {len(X_aug)}")

    # 2. 训练模型
    # model = StockPricePredictor(input_size=6, hidden_size=16)
    # trainer = StockTrainer(model)
    # trainer.train(X_aug, y_aug, scaler_params, epochs=3000)
    #
    # 3. 使用训练好的模型预测
    predictor = StockPricePredictorWrapper()


    for i in range(len(stock_a)-5):
        last_3_days_a = [stock_a[i], stock_a[i+1],stock_a[i+2]]
        last_3_days_b = [stock_b[i], stock_b[i+1],stock_b[i+2]]

        # # 预测次日收盘价
        predicted_price = predictor.predict(last_3_days_a, last_3_days_b)
        actual_next_day = stock_b[i+3]   # 模拟实际值
        #
        print(f"- 预测: {predicted_price:.2f}    实际: {actual_next_day:.2f}")